import math
import numpy as np

from History import *

class Model:
    def __init__(self):
        self.k1 = 0
        self.k2 = 0
        self.k3 = 0
        self.k4 = 0
        self.k5 = 0
        self.k6 = 0
        self.a0 = 0
        self.aDot0 = 0
        
    def getControls(self, dt, a, aDot):
        min = 1e-5
        if self.k1 < min and self.k2 < min and self.k4 < min and self.k5 < min:
            c1 = 0
            c2 = 0
        elif self.k1 < min and self.k2 < min:
            if self.k4 >= min:
                c2 = 0
                c1 = -(self.k5*c2*dt + self.k6*a*dt + self.aDot0 - aDot) / (self.k4*dt)
            else:
                c1 = 0
                c2 = -(self.k4*c1*dt + self.k6*a*dt + self.aDot0 - aDot) / (self.k5*dt)
        elif self.k4 < min and self.k5 < min:
            if self.k2 >= min:
                c1 = 0
                c2 = -(self.k1*c1*dt + self.k3*aDot*dt + self.a0) / (self.k2*dt)
        else:
            c2 = self.k1*(aDot - self.aDot0 - self.k6*a*dt) + self.k4*(self.a0-a+self.k3*aDot*dt)
            c2 = c2 / (dt*(self.k1*self.k5 - self.k4*self.k2))
            if self.k1 < min and self.k4 < min:
                c1 = 0
            elif self.k1 >= min:
                c1 = (a - self.a0 - self.k2*c2*dt - self.k3*aDot*dt) / (self.k1*dt)
            else:
                c1 = (aDot - self.aDot0 - self.k5*c2*dt - self.k6*a*dt) / (self.k4*dt)
        return (c1, c2)
        
    def setParameters(self, e0, e1, e2, e3):
        t1 = e1.time - e0.time
        t2 = e2.time - e0.time
        t3 = e3.time - e0.time
        alpha0 = e0.state.arm1angle
        alphaT1 = e1.state.arm1angle
        alphaT2 = e2.state.arm1angle
        alphaT3 = e3.state.arm1angle
        rhs = np.array([alphaT1 - alpha0, alphaT2 - alpha0, alphaT3 - alpha0])
        a11 = e1.state.arm1control1 * t1
        a12 = e1.state.arm1control2 * t1
        a13 = e1.state.arm1angularVelocity * t1
        a21 = e2.state.arm1control1 * t2
        a22 = e2.state.arm1control2 * t2
        a23 = e2.state.arm1angularVelocity * t2
        a31 = e3.state.arm1control1 * t3
        a32 = e3.state.arm1control2 * t3
        a33 = e3.state.arm1angularVelocity * t3
        a = np.array([[a11,a12,a13], [a21,a22,a23], [a31,a32,a33]])
        #x = np.linalg.solve(a, rhs)
        x = np.linalg.lstsq(a, rhs)
        print x
        alphaDot0 = e0.state.arm1angularVelocity
        alphaDotT1 = e1.state.arm1angularVelocity
        alphaDotT2 = e2.state.arm1angularVelocity
        alphaDotT3 = e3.state.arm1angularVelocity
        rhs = np.array([alphaDotT1 - alphaDot0, alphaDotT2 - alphaDot0, alphaDotT3 - alphaDot0])
        a11 = e1.state.arm1control1 * t1
        a12 = e1.state.arm1control2 * t1
        a13 = e1.state.arm1angle * t1
        a21 = e2.state.arm1control1 * t2
        a22 = e2.state.arm1control2 * t2
        a23 = e2.state.arm1angle * t2
        a31 = e3.state.arm1control1 * t3
        a32 = e3.state.arm1control2 * t3
        a33 = e3.state.arm1angle * t3
        a = np.array([[a11,a12,a13], [a21,a22,a23], [a31,a32,a33]])
        #x = np.linalg.solve(a, rhs)
        x = np.linalg.lstsq(a, rhs)
        print x
    
    def setParameters2(self, e1, e2, e3):
        t = ( e1.time + e2.time + e3.time ) / 3.0
        s0 = State()
        s0.arm1angle = ( e1.state.arm1angle + e2.state.arm1angle + e3.state.arm1angle ) / 3.0
        s0.arm1angularVelocity = ( e1.state.arm1angularVelocity + e2.state.arm1angularVelocity + e3.state.arm1angularVelocity ) / 3.0
        s0.arm1control1 = ( e1.state.arm1control1 + e2.state.arm1control1 + e3.state.arm1control1 ) / 3.0
        s0.arm1control2 = ( e1.state.arm1control2 + e2.state.arm1control2 + e3.state.arm1control2 ) / 3.0
        e0 = Event(s0, t)
        
        t1 = e1.time - e0.time
        t2 = e2.time - e0.time
        t3 = e3.time - e0.time
        alpha0 = e0.state.arm1angle
        alphaT1 = e1.state.arm1angle
        alphaT2 = e2.state.arm1angle
        alphaT3 = e3.state.arm1angle
        rhs = np.array([alphaT1 - alpha0, alphaT2 - alpha0, alphaT3 - alpha0])
        a11 = e1.state.arm1control1 * t1
        a12 = e1.state.arm1control2 * t1
        a13 = e1.state.arm1angularVelocity * t1
        a21 = e2.state.arm1control1 * t2
        a22 = e2.state.arm1control2 * t2
        a23 = e2.state.arm1angularVelocity * t2
        a31 = e3.state.arm1control1 * t3
        a32 = e3.state.arm1control2 * t3
        a33 = e3.state.arm1angularVelocity * t3
        a = np.array([[a11,a12,a13], [a21,a22,a23], [a31,a32,a33]])
        #x = np.linalg.solve(a, rhs)
        x = np.linalg.lstsq(a, rhs)
        print x
        alphaDot0 = e0.state.arm1angularVelocity
        alphaDotT1 = e1.state.arm1angularVelocity
        alphaDotT2 = e2.state.arm1angularVelocity
        alphaDotT3 = e3.state.arm1angularVelocity
        rhs = np.array([alphaDotT1 - alphaDot0, alphaDotT2 - alphaDot0, alphaDotT3 - alphaDot0])
        a11 = e1.state.arm1control1 * t1
        a12 = e1.state.arm1control2 * t1
        a13 = e1.state.arm1angle * t1
        a21 = e2.state.arm1control1 * t2
        a22 = e2.state.arm1control2 * t2
        a23 = e2.state.arm1angle * t2
        a31 = e3.state.arm1control1 * t3
        a32 = e3.state.arm1control2 * t3
        a33 = e3.state.arm1angle * t3
        a = np.array([[a11,a12,a13], [a21,a22,a23], [a31,a32,a33]])
        #x = np.linalg.solve(a, rhs)
        x = np.linalg.lstsq(a, rhs)
        print x

    def printMe(self):
        print "k1 = " + str(self.k1)
        print "k2 = " + str(self.k2)
        print "k3 = " + str(self.k3)
        print "k4 = " + str(self.k4)
        print "k5 = " + str(self.k5)
        print "k6 = " + str(self.k6)
        print "a0 = " + str(self.a0)
        print "aD0= " + str(self.aDot0)
